import os
import sqlite3
from pathlib import Path
import argparse
from datetime import datetime

def check_and_connect_db(db_path):
    """检查并连接到SQLite数据库"""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    # 确保表存在
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS zip_hash (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            file_path TEXT NOT NULL,
            file_name TEXT NOT NULL,
            file_size INTEGER NOT NULL,
            create_time TEXT NOT NULL
        )
    ''')
    # 清空表数据
    cursor.execute('DELETE FROM zip_hash')
    conn.commit()
    return conn, cursor

def insert_zip_info(conn, cursor, file_path, file_name, file_size, create_time):
    """插入ZIP文件信息到数据库"""
    cursor.execute('''
        INSERT INTO zip_hash (file_path, file_name, file_size, create_time)
        VALUES (?, ?, ?, ?)
    ''', (file_path, file_name, file_size, create_time))
    conn.commit()

def scan_and_insert_zip_info(db_directory, db_conn, cursor):
    """扫描指定目录及其子目录下的所有ZIP文件，并插入数据库"""
    db_dir_path = Path(db_directory)

    if not db_dir_path.is_dir():
        print(f"错误: {db_directory} 不是一个有效的目录")
        return

    # 使用rglob递归查找目录下所有的ZIP文件
    zip_files = list(db_dir_path.rglob('*.zip'))

    if not zip_files:
        print("在目录及其子目录中没有找到任何ZIP文件")
        return

    for zip_file in zip_files:
        try:
            file_size = zip_file.stat().st_size
            relative_path = str(zip_file.relative_to(db_dir_path).parent)  # 获取相对路径但不包括文件名
            file_name = zip_file.name  # 单独获取文件名
            create_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

            insert_zip_info(db_conn, cursor, relative_path, file_name, file_size, create_time)
            print(f"已插入记录: 文件路径 {relative_path}, 文件名 {file_name}, 文件大小 {file_size}")
        except Exception as e:
            print(f"处理文件 {zip_file} 时出错: {e}")

def main(db_directory):
    db_path = str(Path(db_directory) / 'file_hashes.db')
    conn, cursor = check_and_connect_db(db_path)

    scan_and_insert_zip_info(db_directory, conn, cursor)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="扫描指定目录及其所有子目录下的所有ZIP文件，并插入到SQLite数据库")
    parser.add_argument('--db', required=True, help="要扫描的目录路径和SQLite数据库存放位置")

    args = parser.parse_args()

    db_directory = args.db

    main(db_directory)